import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter

from utils_vasp.band import get_band
from utils_vasp.dos import get_dos

# get_band() returns distance, band-energy, label-position, and occupation.
dist, band_energy, label_position, occ = get_band(E_fermi=6.1650)

# get_dos() returns energy, total_dos, and partial_dos
dos_energy, dos_total, dos_partial = get_dos("../calc_dos/DOSCAR")

# Sets the size of figure.
F = plt.figure(figsize=(5,3))  # figure's size is in inches

# Need to draw two subfigures. Specifies region of the first one. 
ax_band = plt.axes([0.15, 0.15, 0.60, 0.80])

# Draw vertical lines using a slightly thinner line.
# These lines should be drawn below band lines.
k_label = [r"$\Gamma$", r"$M$", r"$L$", r"$A$", r"$\Gamma$", r"$K$", r"$M$"]
for x in label_position :
    plt.axvline(x, linewidth=0.75)

# Draws the band curves.
for i in range(len(band_energy)):
    plt.plot(dist, band_energy[i])

# Sets the labels of xticks.
plt.xticks(label_position, k_label)

# Matplotlib enlarges slightly the shown range. Shows exactly the range.
plt.xlim(label_position[0], label_position[-1])

# Sets y-axis's label and range.
plt.ylabel(r"$E-E_f$ (eV)")
plt.ylim(-16.0, 11.0)

# Sets the range of the DOS subfigure.
ax_dos = plt.axes([0.80, 0.15, 0.15, 0.80])

# Draws the dos. Note the x-axis is dos, and y-axis is energy.
plt.plot(dos_total, dos_energy, 'k-')

# Hides the numbers because they are shown in the band subplot.
ax_dos.yaxis.set_major_formatter(NullFormatter() )

# Sets y-axis's range the same as that of the band subplot.
plt.ylim(-16.0, 11.0)

# Sets x-axis's range and label.
plt.xlim(0, 5)
plt.xlabel(r"DOS (eV$^{-1}\,$Cell$^{-1}$)")

# Saves the figure into a file (overwritten) and shows it on screen.
plt.savefig("band_dos.png", dpi=300)
plt.show()
